import time
import mujoco
import mujoco.viewer
import numpy as np

from SDK import SyncTeacherConnection, RealMachineConnection


# def on_joint_angles(degrees):
#     # print('Received joint angles: {}'.format(degrees))
#     global target_positions
#     target_positions = np.deg2rad(degrees)
#     # print('target_positions: {}'.format(target_positions))


# 回放关节角的数据
# 回放关节角度数据
def replay_joint_angles(filename='robot_movement.npy'):
    # 从文件中加载角度数据
    with open(filename, 'rb') as f:
        angles_list = []
        while True:
            try:
                angles = np.load(f)
                angles_list.append(angles)
            except:
                break
    return angles_list

if __name__ == "__main__":
    # 加载示教器连接
    teacher = SyncTeacherConnection('config.yml')
    # 加载真实机器人连接
    machine = RealMachineConnection('config.yml')
    
    # Load the model
    model = mujoco.MjModel.from_xml_path('asserts/lift_cube2.xml')
    data = mujoco.MjData(model)

    # 初始化目标位置
    target_positions = np.zeros(6)

    # 回放文件名
    replay_filename = 'robot_movement.npy'

    # 加载录制的关节角度数据
    recorded_angles = replay_joint_angles(replay_filename)

    # 回放模式
    replay_index = 0
    
    # 启动查看器
    with mujoco.viewer.launch_passive(model, data,show_left_ui=False, show_right_ui=False) as viewer:
        # 隐藏左边的控制按钮
        with viewer.lock():
            viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_JOINT] = 0
            viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_ACTUATOR] = 0
            viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_CAMERA] = 0
            viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_LIGHT] = 0
            viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_CONSTRAINT] = 0
            viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_CONTACTFORCE] = 0
            viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = 0
            viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_COM] = 0
            viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_TRANSPARENT] = 0
            
            viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_TENDON] = 0
            viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_RANGEFINDER] = 0
            viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_CONVEXHULL] = 0
        
        while viewer.is_running() :
            step_start = time.time()

            # conn.request_joint_angles(on_joint_angles)
            # angles = teacher.request_joint_angles()
            if replay_index < len(recorded_angles):
            # 获取录制的关节角度
                angles = recorded_angles[replay_index]
                replay_index += 1

                # 设置真实机器人的关节角度
                machine.set_joint_angles(angles)

                # 设置仿真模型的关节角度
                degrees = list(angles)
                degrees[5] = 90 - degrees[5]
                data.ctrl[:] = np.deg2rad(degrees)


           
            mujoco.mj_step(model, data)

            # 同步查看器
            viewer.sync()
            # time.sleep(0.01)